#!/usr/bin/python3

import subprocess
import os
import sys
import argparse
import matplotlib.pyplot as plt
import matplotlib.ticker as tick
import math
from cycler import cycler



# LB/CC mode matching
cc_modes = {
    1: "dcqcn",
    3: "hp",
    7: "timely",
    8: "dctcp",
}
lb_modes = {
    0: "fECMP",
    2: "DRILL",
    3: "conga",
    6: "letflow",
    9: "conweave",
    13: "lessflow",
}
topo2bdp = {
    "leaf_spine_128_100G_OS2": 104000,  # 2-tier
    "fat_k4_100G_OS2": 153000, # 3-tier -> core 400G
}


C = [
    'xkcd:grass green',
    'xkcd:blue',
    'xkcd:purple',
    'xkcd:orange',
    'xkcd:tomato',
    'xkcd:brick red',
    'xkcd:black',
    'xkcd:teal',
    'xkcd:brown',
    'xkcd:grey',
]

LS = [
    'solid',
    'dashed',
    'dotted',
    'dashdot'
]

M = [
    'o',
    's',
    'x',
    'v',
    'D'
]

H = [
    '//',
    'o',
    '***',
    'x',
    'xxx',
]

def setup():
    """Called before every plot_ function"""

    def lcm(a, b):
        return abs(a*b) // math.gcd(a, b)

    def a(c1, c2):
        """Add cyclers with lcm."""
        l = lcm(len(c1), len(c2))
        c1 = c1 * (l//len(c1))
        c2 = c2 * (l//len(c2))
        return c1 + c2

    def add(*cyclers):
        s = None
        for c in cyclers:
            if s is None:
                s = c
            else:
                s = a(s, c)
        return s

    plt.rc('axes', prop_cycle=(add(cycler(color=C),
                                   cycler(linestyle=LS),
                                   cycler(marker=M))))
    plt.rc('lines', markersize=5)
    plt.rc('legend', handlelength=3, handleheight=1.5, labelspacing=0.25)
    plt.rcParams["font.family"] = "sans"
    plt.rcParams["font.size"] = 16
    plt.rcParams['pdf.fonttype'] = 42
    plt.rcParams['ps.fonttype'] = 42


def getFilePath():
    dir_path = os.path.dirname(os.path.realpath(__file__))
    print("File directory: {}".format(dir_path))
    return dir_path


def get_rate_data(filename):
    res = {}     # key是时间，value是该时间下的总字节数
    accumulate_data = {}
    with open(filename, 'r') as f:
        for line in f:
            parsed_line = line.replace("\n", "").split(",")
            now_ts = int(parsed_line[0])
            now_swid = int(parsed_line[1])
            now_portid = int(parsed_line[2])
            now_val = int(parsed_line[3])
            
            # if now_ts < time_start or now_ts > time_end:
            #     continue
            
           
            if now_ts not in accumulate_data:
                accumulate_data[now_ts] = now_val
            else:
                accumulate_data[now_ts] += now_val
    
    sorted_keys = sorted(accumulate_data.keys())
    cnt = 0
    prev = sorted_keys[0]
    for key in sorted_keys:
        if cnt > 0:
            res[key] = (accumulate_data[key] - accumulate_data[prev]) / (key - prev)
            prev = key
        cnt += 1
                
    return res


def main():
    file_dir = getFilePath()
    fig_dir = file_dir + "/figures"
    output_dir = file_dir + "/../mix_k2flow1G/output"
    history_filename = file_dir + "/../mix_k2flow1G/.history"
    
    
    com_dict = {
               # lb_mode, enabled_irn, split_num, flow_size
        'ECMP': ('0', '0', '1', '100000000'),
        'DRILL': ('2', '0', '1', '100000000'),
        'ConWeave': ('9', '0', '1', '100000000'),
        'LetFlow': ('6', '0', '1', '100000000'),
        'ParaLet(m=8)': ('13', '0', '8', '100000000'),
        'ParaLet(m=16)': ('13', '0', '16', '100000000'),
        'ParaLet(m=32)': ('13', '0', '32', '100000000'),
    }
    need_com_par = com_dict.values()
    swapped_dict = dict(zip(com_dict.values(), com_dict.keys()))
    
    
    plot_dict = {}
    for key in com_dict.keys():
        plot_dict[key] = None


    
    # {simulday},{config_ID},{cc_mode},{lb_mode},{enabled_pfc},{enabled_irn},{topo},{split_num},{flow_size}\n
    with open(history_filename, 'r') as hf:
        row_str_list = hf.read().splitlines()
        for row_str in row_str_list:
            row = row_str.split(',')
            cur_par = (row[3], row[5], row[7], row[8])
            if cur_par in need_com_par:
                exp_file_path = output_dir + f'/{row[1]}'
                uplink_file = exp_file_path + f'/{row[1]}_out_uplink.txt'
                res = get_rate_data(uplink_file)
                plot_dict[swapped_dict[cur_par]] = res

    # 设置图形大小
    plt.figure(figsize=(5, 3))

    # 采样频率，例如每隔100个点取一个点
    sampling_rate = 50
    
    # 遍历数据字典，绘制每个方法的数据点
    for method, data in plot_dict.items():
        # 分别获取横坐标和纵坐标的值
        x_values = list(data.keys())
        y_values = list(data.values())
        
        # 进行线性采样
        sampled_x_values = x_values[::sampling_rate]
        sampled_y_values = y_values[::sampling_rate]
        
        
        # 横坐标单位换算，ns变ms，且从退到从0开始
        sampled_x_values = [x / 1000000 - 2000 for x in sampled_x_values]
        # 纵坐标单位换算，(Byte / ns) 变成Gbps
        sampled_y_values = [y * 8 / 100 for y in sampled_y_values]
        
        print("Max total throughput,", method, max(sampled_y_values))
        
        # 绘制方法的数据点
        plt.plot(sampled_x_values, sampled_y_values, markersize = 0, linewidth = 3.0, label = method)

    # 添加图例
    plt.legend()


    # 添加标题和轴标签
    plt.xlabel('Time (ms)')
    plt.ylabel('TT (x100 Gbps)')

    # 显示网格
    # plt.grid(True)
    plt.legend(frameon=False, fontsize=12, facecolor='white')
    plt.grid(which='minor', alpha=0.2)
    plt.grid(which='major', alpha=0.5)
    fig_name = os.path.dirname(os.path.abspath(__file__)) + '/' +  __file__.split('/')[-1].split('.')[0] + '.pdf'
    print(fig_name)
    plt.savefig(f"{fig_name}", bbox_inches='tight')


if __name__=="__main__":
    setup()
    main()
